{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Using FlexGen to Offload OPT Models' Weights to RAM and Disk\n",
        "This notebook is a companion of chapter 8 of the \"Domain Specific LLMs in Action\" book, author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2024.  \n",
        "The code in this notebook is to perform inference with a Meta AI's OPT models, by offloading part of models' weights from VRAM to RAM and/or disk, using the [FlexGen](https://github.com/FMInference/FlexLLMGen/) generation engine programmatically. While the code refers to the [OPT 1.3 B](https://huggingface.co/facebook/opt-1.3b) model, the same applies to any other model from the same family. It requires hardware acceleration to be executed.  \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "RHR7quvS5Ld1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the FlexGen from source."
      ],
      "metadata": {
        "id": "cXnjUWSo65zl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/FMInference/FlexLLMGen.git\n",
        "%cd FlexLLMGen\n",
        "!pip install -e ."
      ],
      "metadata": {
        "id": "PpVXlrLZFtG6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Using FlexGen and the Transformers library programmatically"
      ],
      "metadata": {
        "id": "IjHOD3HNnwx0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import the required FlexGen classes."
      ],
      "metadata": {
        "id": "U5I5V7gA72Tl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from flexllmgen.flex_opt import (Policy, OptLM, ExecutionEnv, CompressionConfig,\n",
        "        str2bool)"
      ],
      "metadata": {
        "id": "ZpJyirZsoA3B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the OPT 1.3 B tokenizer form the Hugging Face's Hub."
      ],
      "metadata": {
        "id": "Ecdzz6_h8AD1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "model_id = \"facebook/opt-1.3b\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"left\")\n",
        "tokenizer.add_bos_token = False\n",
        "stop = tokenizer(\"\\n\").input_ids[0]"
      ],
      "metadata": {
        "id": "9KRpa3tJ6knM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Setup the FlexGen execution environment."
      ],
      "metadata": {
        "id": "Rtmal7Oc8LcT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "offload_dir = './flexgen_offload'\n",
        "env = ExecutionEnv.create(offload_dir)"
      ],
      "metadata": {
        "id": "LpT-1eLU7cWq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Prepare a list of prompts for batch inference."
      ],
      "metadata": {
        "id": "Abl4xDU98jPC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompts = [\n",
        "    \"Question: Where were the 2004 Olympics held?\\n\"\n",
        "    \"Answer: Athens, Greece\\n\"\n",
        "    \"Question: What is the longest river on the earth?\\n\"\n",
        "    \"Answer:\",\n",
        "\n",
        "    \"Extract the airport codes from this text.\\n\"\n",
        "    \"Text: \\\"I want a flight from New York to San Francisco.\\\"\\n\"\n",
        "    \"Airport codes: JFK, SFO.\\n\"\n",
        "    \"Text: \\\"I want you to book a flight from Phoenix to Las Vegas.\\\"\\n\"\n",
        "    \"Airport codes:\",\n",
        "]"
      ],
      "metadata": {
        "id": "Fwm6oye_8LUs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Setup an offloading policy."
      ],
      "metadata": {
        "id": "W9HOysuq8mqg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "policy = Policy(len(prompts), 1,\n",
        "                70, 30, 70, 30, 100, 0,\n",
        "                overlap=True, sep_layer=True, pin_weight=True,\n",
        "                cpu_cache_compute=True, attn_sparsity=1.0,\n",
        "                compress_weight=True,\n",
        "                comp_weight_config=CompressionConfig(\n",
        "                    num_bits=4, group_size=64,\n",
        "                    group_dim=0, symmetric=False),\n",
        "                compress_cache=False, # Set compress_cache to False\n",
        "                comp_cache_config=CompressionConfig(\n",
        "                    num_bits=4, group_size=64,\n",
        "                    group_dim=2, symmetric=False)\n",
        "                )"
      ],
      "metadata": {
        "id": "cNMUhRWC8FZr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Prepare the model to be executed through the FlexGen inference engine and following the preliminary defined offloading policies. This step also downloads the model's checkpoints from the Hugging Face's Hub and manages the conversion process."
      ],
      "metadata": {
        "id": "gvOFzEh_8v_K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "path = '~/opt_weights'\n",
        "model = OptLM(model_id, env, path, policy)"
      ],
      "metadata": {
        "id": "YPPXnIrh6zgU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate text for the given set of prompts and then display the generated result for each one."
      ],
      "metadata": {
        "id": "t5gUNQNN9Bhs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Generate...\")\n",
        "inputs = tokenizer(prompts, padding=\"max_length\", max_length=128)\n",
        "output_ids = model.generate(\n",
        "    inputs.input_ids,\n",
        "    do_sample=True,\n",
        "    temperature=0.7,\n",
        "    max_new_tokens=32,\n",
        "    stop=stop)\n",
        "outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)\n",
        "print(\"Outputs:\\n\" + 70 * '-')\n",
        "for i in [0, len(outputs)-1]:\n",
        "    print(f\"{i}: {outputs[i]}\")\n",
        "    print(\"-\" * 70)"
      ],
      "metadata": {
        "id": "-eFB8TIF-JBN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Shutdown the FlexGen execution environment when done."
      ],
      "metadata": {
        "id": "6W9F4pEI9LDM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Shutting down...\")\n",
        "env.close_copy_threads()"
      ],
      "metadata": {
        "id": "EqCEdaED-QhD"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}